import json
import torch
import torch.nn as nn
from functools import partial
from mmengine import Config
from mmdet.registry import MODELS
from mmengine.runner import Runner
from mmdet.utils import register_all_modules
from mmengine.model import revert_sync_batchnorm
from torchinfo import summary
register_all_modules()

class ModelInspector:
    def __init__(self):
        self.config_path = None
        self.config = None
        self.model = None
        self.input_size = None

    def inspect(self, config_path, dataset_cfg_dict):
        self.config_path = config_path
        self.config = Config.fromfile(config_path)
        if dataset_cfg_dict:
            self.config.merge_from_dict(dataset_cfg_dict)
        self.model = None
        self.input_size = None

        # 开始加载对应的信息
        self.setup_model()
        data = self.prepare_data_sample()
        self.modify_forward(data)
        json_summary = self.generate_summary()
        return json_summary

    def setup_model(self):
        self.model = MODELS.build(self.config.model)
        if torch.cuda.is_available():
            self.model = self.model.cuda()
        self.model.eval()

    def prepare_data_sample(self):
        data_loader = Runner.build_dataloader(self.config.val_dataloader)
        data_iter = iter(data_loader)
        data_batch = next(data_iter)
        data = self.model.data_preprocessor(data_batch)
        self.input_size = data['inputs'].shape
        return data

    def modify_forward(self, data):
        _forward = self.model.forward
        self.model.forward = partial(_forward, data_samples=data['data_samples'])

    def generate_summary(self):
        model_summary = summary(self.model, input_size=self.input_size)
        layers_info = self.parse_layers(model_summary.summary_list)
        summary_dict = {
            "total_params": model_summary.total_params,
            "trainable_params": model_summary.trainable_params,
            "layers": layers_info
        }
        return json.dumps(summary_dict, indent=2)

    def parse_layers(self, layers):
        parsed_layers = []
        for layer in layers:
            parsed_layer = {
                "class_name": layer.class_name,
                "var_name": layer.var_name,
                "depth": layer.depth,
                "depth_index": layer.depth_index,
                "input_size": str(layer.input_size),
                "output_size": str(layer.output_size),
                "num_params": layer.num_params if not layer.is_recursive else "recursive",
                "trainable_params": layer.trainable_params,
                "children": self.parse_layers(layer.children) if layer.children else []
            }
            parsed_layers.append(parsed_layer)
        return parsed_layers

if __name__ == '__main__':
    config_file = '/home/lixiang/.config/labelapp/output/YYWahrFLTiiQskmIS9REg/faster-rcnn_r50-caffe_fpn_ms-2x_coco.py'
    inspector = ModelInspector()
    json_summary = inspector.inspect(config_file)
